#!/Pub/Users/fuxj/.conda/envs/R_clusterProfiler/bin/Rscript
suppressPackageStartupMessages(library(argparse))


GSEA_cluster <- function(
    deg_res = NULL, pvalueCutoff = 0.05,
    pAdjustMethod = "none", seed = 1110, saveplot = F, output_dir = NULL, var_name = NULL) {
  suppressPackageStartupMessages(library(RColorBrewer))
  suppressPackageStartupMessages(library(org.Hs.eg.db))
  suppressPackageStartupMessages(library(clusterProfiler))
  suppressPackageStartupMessages(library(cowplot))
  suppressPackageStartupMessages(library(enrichplot))
  suppressPackageStartupMessages(library(tidyverse))

  # browser()

  tryCatch(stopifnot(exprs = {
    "log2FC" %in% colnames(deg_res)
    "gene" %in% colnames(deg_res)
  }), error = \(e) {
    print(e)
    simpleMessage("file_of_deg需要包含gene与log2FC两列。")
  })

  deg_all <- deg_res %>%
    arrange(desc(log2FC)) %>%
    dplyr::select(gene, log2FC)

  if (is.null(var_name)) {
    var_name <- paste0(sample(letters, 4), collapse = "")
  }

  gene_name <- bitr(deg_all$gene, fromType = "SYMBOL", toType = "ENTREZID", OrgDb = "org.Hs.eg.db", drop = F)

  tmp_data <- left_join(gene_name, deg_all %>% dplyr::rename(SYMBOL = gene))
  tmp_data <- na.omit(tmp_data)

  genelist <- tmp_data$log2FC
  names(genelist) <- tmp_data$ENTREZID

  gsea_kegg <- gseKEGG(
    geneList = genelist,
    verbose = T,
    seed = F,
    # nPerm = 1000,#
    keyType = "kegg", # 可以选择"kegg",'ncbi-geneid', 'ncib-proteinid' and 'uniprot'
    organism = "hsa", # 定义物种,
    pvalueCutoff = pvalueCutoff, # 自定义pvalue的范围
    pAdjustMethod = pAdjustMethod # 校正p值的方法
  )

  gsea_GO <- gseGO(
    geneList = genelist,
    ont = "all",
    OrgDb = org.Hs.eg.db,
    keyType = "ENTREZID",
    pvalueCutoff = pvalueCutoff,
    pAdjustMethod = pAdjustMethod,
    verbose = TRUE,
    seed = F
  )

  gsea_kegg@result <- gsea_kegg@result %>% mutate(Description = str_wrap(Description,45))
  gsea_GO@result <- gsea_GO@result %>% mutate(Description = str_wrap(Description,45))

  gsea_kegg <- gsea_kegg %>%
    arrange(desc(enrichmentScore))

  gsea_res_list <- list("gsea_kegg" = gsea_kegg, "gsea_GO" = gsea_GO)

  p_gsea_kegg_line <- gseaplot2(
    x = gsea_kegg,
    geneSetID = gsea_kegg$ID[1:4], # 只显示前4个GSEA的结果
    title = "GESA KEGG", # 标题
    color = brewer.pal(8, "Set1"), # 颜色
    pvalue_table = FALSE,
    ES_geom = "line"
  ) # enrichment scored的展现方式 'line' or 'dot')

  p_gsea_kegg_line_tail <- gseaplot2(
    x = gsea_kegg,
    geneSetID = tail(gsea_kegg$ID, 4), # 只显示后4个GSEA的结果
    title = "GESA KEGG", # 标题
    color = brewer.pal(8, "Set1"), # 颜色
    pvalue_table = FALSE,
    ES_geom = "line"
  ) # enrichment scored的展现方式 'line' or 'dot')


  gsea_GO <- gsea_GO %>%
    arrange(desc(enrichmentScore))

  p_gsea_GO_line <- gseaplot2(
    x = gsea_GO,
    geneSetID = head(gsea_GO$ID, 4), # 只显示后4个GSEA的结果
    title = "GESA GO", # 标题
    color = brewer.pal(8, "Set1"), # 颜色
    pvalue_table = FALSE,
    ES_geom = "line"
  ) # enrichment scored的展现方式 'line' or 'dot')

  p_gsea_GO_line_tail <- gseaplot2(
    x = gsea_GO,
    geneSetID = tail(gsea_GO$ID, 4), # 只显示后4个GSEA的结果
    title = "GESA GO", # 标题
    color = brewer.pal(8, "Set1"), # 颜色
    pvalue_table = FALSE,
    ES_geom = "line"
  ) # enrichment scored的展现方式 'line' or 'dot')


  gsea_line <- cowplot::plot_grid(p_gsea_kegg_line, p_gsea_GO_line,
    ncol = 2,
    nrow = 1
  )

  x <- list(
    p_gsea_kegg_line, p_gsea_GO_line
  )

  names(x) <- c(
    "p_gsea_kegg_line", "p_gsea_GO_line"
  )


  if (saveplot) {
    if (!dir.exists(sprintf("%s/GSEA_cluster_%s", output_dir, var_name))) {
      dir.create(sprintf("%s/GSEA_cluster_%s", output_dir, var_name), recursive = T)
    } else {
      message(sprintf("Dir '%s/GSEA_cluster_%s' is existed.", output_dir, var_name))
    }

    dir_now <- sprintf("%s/GSEA_cluster_%s/", output_dir, var_name)

    gsea_kegg@result %>% write_delim(
      x = ., file = sprintf("%sGSEA_KEGG_%s.txt", dir_now, var_name),
      delim = "\t", quote = "none"
    )

    gsea_GO@result %>% write_delim(
      x = ., file = sprintf("%sGSEA_GO_%s.txt", dir_now, var_name),
      delim = "\t", quote = "none"
    )

    c("p_gsea_kegg_line", "p_gsea_kegg_line_tail", "p_gsea_GO_line", "p_gsea_GO_line_tail") %>%
      walk(~ {
        pdf(file = sprintf("%sgsea_line_%s.pdf", dir_now, .x), width = 6, height = 4.8)
        print(get0(.x))
        dev.off()
      })

    saveRDS(gsea_res_list, sprintf("%sgsea_res_list_%s.rds", dir_now, var_name))
  }

  return(x)
}

wrp <- "直接来利用小俊姐的conda富集分析环境，对指定基因进行GO, KEGG富集分析。tail字样代表另一个方向的富集结果。"
parser <- ArgumentParser(
    description = wrp,
    prog = "gsea_cmd.r",
    formatter_class = "argparse.RawTextHelpFormatter"
)

parser$add_argument("-od", "--output_dir",
    action = "store",
    required  = TRUE,
    default = "./", metavar = "",
    help = "代表结果输出目录 ,默认当前命令执行目录，`./`"
)
parser$add_argument("-f", "--file_of_deg",
    action = "store",
    required  = TRUE,
    default = "/Pub/Users/wangyk/project/Poroject/F240614003_m6A_typ2_source_code/test/10/4./deg_res_nmf_cls.txt", metavar = "",
    help = "基因文件所在路径。要求第一列为gene，利用第一列的基因进行分析。"
)
parser$add_argument("-v", "--var_name",
    action = "store",
    default = "_", metavar = "",
    help = "保存文件中识别的关键字，默认为下划线"
)

parser$add_argument("-p", "--pvalue",
    action = "store",type = 'numeric',
    default = .05, metavar = "",
    help = "结果筛选P值，默认0.05"
)
parser$add_argument("-s", "--symbol",
    default = 'gene', metavar = "",required  = TRUE,
    help = "基因symbol所在列列名，默认为`gene`"
)
parser$add_argument("-l", "--log2FC",
    action = "store",required  = TRUE,
    default = "log2FC", metavar = "",
    help = "log2FC所在列列名，默认为`log2FC`"
)
parser$add_argument("-a", "--adj_p_method",
    action = "store",
    required  = TRUE,
    default = "BH", metavar = "",
    help = "p值矫正方法，默认BH"
)


args_in_process <- parser$parse_args()

# od check
if (dir.exists(args_in_process$output_dir)) {
   dir.create(args_in_process$output_dir, showWarnings = F,recursive = T)
}

# file check
.file_of_deg_check <- function() {
  tryCatch(stopifnot(file.exists(args_in_process$file_of_deg)),
    error = function(e) {
      cli::cli_alert_danger( "指定的基因文件不存在。{.path {args_in_process$file_of_deg}}")
    }
  )
}
.file_of_deg_check()

suppressPackageStartupMessages(library(tidyverse))
# read data
if (str_detect(basename(args_in_process$file_of_deg), "csv")) {
    degs <- read.csv(args_in_process$file_of_deg)
} else {
    degs <- read.delim(args_in_process$file_of_deg)
}

colnames(degs)[which(colnames(degs) == args_in_process$symbol)] <- 'gene'
colnames(degs)[which(colnames(degs) == args_in_process$log2FC)] <- 'log2FC'

main <- \(){
  a <- GSEA_cluster(
    deg_res = degs,
    pvalueCutoff = args_in_process$pvalue,
    pAdjustMethod = args_in_process$adj_p_method,
    seed = 1110,
    saveplot = T,
    output_dir = paste0(args_in_process$output_dir),
    var_name = args_in_process$var_name
  )
}

try(main())



